#include <fstream>
#include <random>

#include <bthread/bthread.h>
#include <glog/logging.h>
#include <gtest/gtest.h>

#include "util/snapshot_vector/snapshot_vector.h"
#include "util/pool/pool.h"

using namespace hawking::indexlib;
DEFINE_int32(snapshot_test_tasks, 128, "num of threads or bthreads");

void* Write(void* arg) {
    SnapshotVector<uint32_t>* snapshot_vec =
        reinterpret_cast<SnapshotVector<uint32_t>*>(arg);
    while (1) {
        snapshot_vec->PushBack(0);
        bthread_usleep(1 * 1000);  // 10 ms
        LOG(INFO) << "cur size " << snapshot_vec->Size();
    }
}

void* Read(void* arg) {
    SnapshotVector<uint32_t>* snapshot_vec =
        reinterpret_cast<SnapshotVector<uint32_t>*>(arg);
    std::unique_ptr<uint32_t> dummy = std::make_unique<uint32_t>();
    while (1) {
        SnapshotVector<uint32_t>::Iterator iter = snapshot_vec->Snapshot();
        while (iter.HasNext()) {
            *dummy = iter.Next();
        }
    }
}

TEST(snapshot_vector_test, one_write_multi_read) {
    SnapshotVector<uint32_t> snapshot_vec;

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, Write, &snapshot_vec);
    sleep(1);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_snapshot_test_tasks);
    for (size_t idx = 0; idx < FLAGS_snapshot_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, Read, &snapshot_vec);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
}

TEST(snapshot_vector_test, one_write_multi_read_use_pool) {
    std::unique_ptr<MMapAllocator> allocator = std::make_unique<MMapAllocator>();
    std::unique_ptr<Pool> pool = std::make_unique<Pool>(allocator.get(), 64);
    SnapshotVector<uint32_t, pool_allocator<uint32_t>> snapshot_vec(pool.get());

    // write bthread
    bthread_t write_bid;
    bthread_start_background(&write_bid, nullptr, Write, &snapshot_vec);
    sleep(1);

    // read bhtreads
    std::vector<bthread_t> read_bids;
    read_bids.reserve(FLAGS_snapshot_test_tasks);
    for (size_t idx = 0; idx < FLAGS_snapshot_test_tasks; ++idx) {
        read_bids.emplace_back();
        bthread_start_background(&read_bids.back(), nullptr, Read, &snapshot_vec);
    }

    bthread_join(write_bid, nullptr);
    for (auto read_bid : read_bids) {
        bthread_join(read_bid, nullptr);
    }
}

int main(int argc, char* argv[]) {
    ::testing::InitGoogleTest(&argc, argv);
    ::google::ParseCommandLineFlags(&argc, &argv, true);
    srand(static_cast<unsigned>(time(0)));
    return RUN_ALL_TESTS();
}
